/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         https://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

#include "itkHoughTransform2DLinesImageFilter.h"
#include "itkThresholdImageFilter.h"
#include "itkGradientMagnitudeImageFilter.h"
#include "itkMath.h"
#include "itkTestingMacros.h"

/**
 * This program looks for straight lines within an image
 * It uses the ITK HoughTransform2DLinesImageFilter.
 * - Read the image.
 * - Apply a gradient and thresholding functions.
 * - Compute the accumulator by running the filter.
 * - Blur the accumulator.
 * - Find maxima in the accumulator.
 * - Display the results
 *
 * It also does unit testing.
 */

namespace
{
bool
Test_GetLines_should_return_empty_list_when_input_image_is_entirely_black()
{
  using PixelType = unsigned char;

  using ImageType = itk::Image<PixelType>;

  using FilterType = itk::HoughTransform2DLinesImageFilter<PixelType, double>;

  // Create a black input image for the filter.
  const auto                    image = ImageType::New();
  constexpr ImageType::SizeType size = { { 32, 32 } };
  image->SetRegions(size);
  image->AllocateInitialized();

  const auto filter = FilterType::New();
  filter->SetInput(image);
  filter->Update();

  if (!filter->GetLines().empty())
  {
    std::cout << "GetLines() should return an empty list when the input image is entirely black." << std::endl;
    return false;
  }
  return true;
}


bool
Test_GetLines_should_return_empty_list_when_NumberOfLines_is_set_to_zero()
{
  using PixelType = unsigned char;

  using ImageType = itk::Image<PixelType>;

  // Create an image.
  const auto image = ImageType::New();
  enum
  {
    sizeX = 32,
    sizeY = 32
  };
  constexpr ImageType::SizeType size = { { sizeX, sizeY } };
  image->SetRegions(size);
  image->AllocateInitialized();

  // Place some line segment in the image.
  for (itk::IndexValueType x = 1; x < (sizeX - 1); ++x)
  {
    const itk::Index<> index = { { x, sizeY / 2 } };
    image->SetPixel(index, 1);
  }

  using FilterType = itk::HoughTransform2DLinesImageFilter<PixelType, double>;

  const auto filter = FilterType::New();

  filter->SetInput(image);
  filter->SetNumberOfLines(0);
  filter->Update();

  // Even when there appears a line segment in the image, GetLines() should return an empty list
  // because SetNumberOfLines(0) was called.
  if (!filter->GetLines().empty())
  {
    std::cout << "GetLines() should return an empty list when NumberOfLines is set to zero." << std::endl;
    return false;
  }
  return true;
}
} // namespace


/** Hough Point structure */
struct HoughPoint
{
  double radius;
  double angle;
};


int
itkHoughTransform2DLinesImageTest(int, char *[])
{
  bool success = true;

  // Define the dimension of the images
  constexpr unsigned int Dimension = 2;

  // Declare the pixel types of the images
  using PixelType = unsigned char;
  using HoughSpacePixelType = double;

  // Declare the types of the images
  using HoughImageType = itk::Image<HoughSpacePixelType, Dimension>;
  using ImageType = itk::Image<PixelType, Dimension>;


  // Create a line image with one line
  auto image = ImageType::New();

  ImageType::RegionType region;

  auto size = ImageType::SizeType::Filled(100);

  ImageType::IndexType index{};

  region.SetSize(size);
  region.SetIndex(index);

  image->SetRegions(region);
  image->AllocateInitialized();

  // Create a line
  constexpr unsigned int lines = 1;
  constexpr double       theta = 0.20; // radians
  constexpr double       radius = 50;

  const double Vx = radius * std::cos(theta);
  const double Vy = radius * std::sin(theta);

  const double norm = std::sqrt(Vx * Vx + Vy * Vy);
  const double VxNorm = Vx / norm;
  const double VyNorm = Vy / norm;

  const unsigned int numberOfPixels = size[0] * size[1];

  for (unsigned int i = 0; i < numberOfPixels; i += 1)
  {
    index[0] = static_cast<long>(Vx - VyNorm * i);
    index[1] = static_cast<long>(Vy + VxNorm * i);

    if (index[0] < static_cast<long>(size[0]) && index[0] >= 0 && index[1] < static_cast<long>(size[1]) &&
        index[1] >= 0)
    {
      image->SetPixel(index, 255);
    }
  }

  // Allocate Hough Space image (accumulator)
  auto m_HoughSpaceImage = HoughImageType::New();
  m_HoughSpaceImage->SetRegions(region);
  m_HoughSpaceImage->Allocate();

  // Apply gradient filter to the input image
  using CastingFilterType = itk::CastImageFilter<ImageType, HoughImageType>;

  auto caster = CastingFilterType::New();
  caster->SetInput(image);


  using GradientFilterType = itk::GradientMagnitudeImageFilter<HoughImageType, HoughImageType>;

  auto gradFilter = GradientFilterType::New();
  gradFilter->SetInput(caster->GetOutput());
  gradFilter->Update();

  /// Apply a threshold to the Grad(InputImage)
  using ThresholdFilterType = itk::ThresholdImageFilter<HoughImageType>;

  auto threshFilter = ThresholdFilterType::New();
  threshFilter->SetInput(gradFilter->GetOutput());
  threshFilter->SetOutsideValue(0);
  constexpr unsigned char lowerThreshold = 10;
  constexpr unsigned char upperThreshold = 200;
  threshFilter->ThresholdOutside(lowerThreshold, upperThreshold);

  threshFilter->Update();

  // Define the HoughTransform filter
  using HoughTransformFilterType = itk::HoughTransform2DLinesImageFilter<HoughSpacePixelType, HoughSpacePixelType>;

  auto houghFilter = HoughTransformFilterType::New();

  ITK_EXERCISE_BASIC_OBJECT_METHODS(houghFilter, HoughTransform2DLinesImageFilter, ImageToImageFilter);


  constexpr float threshold = 2.3;
  houghFilter->SetThreshold(threshold);
  ITK_TEST_SET_GET_VALUE(threshold, houghFilter->GetThreshold());

  constexpr float angleResolution = 200.0;
  houghFilter->SetAngleResolution(angleResolution);
  ITK_TEST_SET_GET_VALUE(angleResolution, houghFilter->GetAngleResolution());

  auto numberOfLines = static_cast<HoughTransformFilterType::LinesListSizeType>(lines);
  houghFilter->SetNumberOfLines(numberOfLines);
  ITK_TEST_SET_GET_VALUE(numberOfLines, houghFilter->GetNumberOfLines());

  constexpr float discRadius = 25.0;
  houghFilter->SetDiscRadius(discRadius);
  ITK_TEST_SET_GET_VALUE(discRadius, houghFilter->GetDiscRadius());

  constexpr float variance = 10;
  houghFilter->SetVariance(variance);
  ITK_TEST_SET_GET_VALUE(variance, houghFilter->GetVariance());


  houghFilter->SetInput(threshFilter->GetOutput());

  houghFilter->Update();

  houghFilter->Simplify();

  const HoughImageType::Pointer accumulator = houghFilter->GetOutput();

  const HoughImageType::ConstPointer simplifyAccumulator = houghFilter->GetSimplifyAccumulator();


  // Blur the accumulator in order to find the maximum
  using GaussianFilterType = itk::DiscreteGaussianImageFilter<HoughImageType, HoughImageType>;

  auto gaussianFilter = GaussianFilterType::New();
  gaussianFilter->SetInput(accumulator);
  double gaussianFilterVariance[Dimension];
  gaussianFilterVariance[0] = variance;
  gaussianFilterVariance[1] = variance;
  gaussianFilter->SetVariance(gaussianFilterVariance);
  gaussianFilter->SetMaximumError(.01f);

  gaussianFilter->Update();

  const HoughImageType::Pointer postProcessImage = gaussianFilter->GetOutput();

  using MinMaxCalculatorType = itk::MinimumMaximumImageCalculator<HoughImageType>;
  auto minMaxCalculator = MinMaxCalculatorType::New();

  const itk::ImageRegionIterator<HoughImageType> it_output(m_HoughSpaceImage,
                                                           m_HoughSpaceImage->GetLargestPossibleRegion());

  itk::ImageRegionIterator<HoughImageType> it_input(postProcessImage, postProcessImage->GetLargestPossibleRegion());


  unsigned int foundLines = 0;

  std::list<HoughPoint> linesList;

  // Search for maxima
  // Each time a maximum is found it is removed by drawing a black disc
  // whose size is defined by the Hough disc radius
  do
  {
    minMaxCalculator->SetImage(postProcessImage);
    minMaxCalculator->ComputeMaximum();
    HoughImageType::PixelType max = minMaxCalculator->GetMaximum();

    for (it_input.GoToBegin(); !it_input.IsAtEnd(); ++it_input)
    {
      if (itk::Math::ExactlyEquals(it_input.Get(), max))
      {
        HoughPoint houghPoint;
        houghPoint.radius = it_input.GetIndex()[0];
        houghPoint.angle =
          ((it_input.GetIndex()[1]) * 2 * itk::Math::pi / houghFilter->GetAngleResolution()) - itk::Math::pi;

        linesList.push_back(houghPoint);

        // Remove a black disc from the Hough space domain
        for (double angle = 0; angle <= 2 * itk::Math::pi; angle += itk::Math::pi / 1000)
        {
          for (double length = 0; length < discRadius; length += 1)
          {
            index[0] = static_cast<long>(it_input.GetIndex()[0] + length * std::cos(angle));
            index[1] = static_cast<long>(it_input.GetIndex()[1] + length * std::sin(angle));
            if (index[0] <= std::sqrt(400.0 * 400 + 400 * 400) && index[0] >= 0 && index[1] <= angleResolution &&
                index[1] >= 0)
            {
              accumulator->SetPixel(index, 0);
            }
          }
        }
        minMaxCalculator->SetImage(accumulator);
        minMaxCalculator->ComputeMaximum();
        max = minMaxCalculator->GetMaximum();

        foundLines++;
        if (foundLines == lines)
        {
          break;
        }
      }
    }
  } while (foundLines < lines);

  // Check the line detection
  auto it_list = linesList.begin();

  constexpr double angleTolerance = 0.1;
  constexpr double radiusTolerance = 1.0;
  while (it_list != linesList.end())
  {
    if (!itk::Math::FloatAlmostEqual(it_list->angle, theta, 10, angleTolerance))
    {
      std::cout << "Failure for line" << std::endl;
      std::cout << "Expected angle: " << theta << ", found: " << it_list->angle << std::endl;
      success = false;
    }
    if (!itk::Math::FloatAlmostEqual(it_list->radius, radius, 10, radiusTolerance))
    {
      std::cout << "Failure for line" << std::endl;
      std::cout << "Expected radius: " << radius << ", found: " << it_list->radius << std::endl;
      success = false;
    }
    else
    {
      std::cout << "Line :"
                << " angle: " << it_list->angle << ", "
                << "radius: " << it_list->radius << std::endl;
    }
    ++it_list;
  }

  success &= Test_GetLines_should_return_empty_list_when_input_image_is_entirely_black();
  success &= Test_GetLines_should_return_empty_list_when_NumberOfLines_is_set_to_zero();

  if (success)
  {
    std::cout << "Test succeeded!" << std::endl;
    return EXIT_SUCCESS;
  }

  std::cout << "Test FAILED!" << std::endl;
  return EXIT_FAILURE;
}
